#include <assert.h>
#include <stdio.h>

#include "../src/mstep.c"

static void
fill (double*           x,
      double          (*func) (void),
      unsigned int      n)
{
  for (unsigned int i = 0; i < n; ++i)
    { 
      x[i] = func ();
    }
}

static void
fill_r (Rating*         x,
        unsigned int    n_workers,
        unsigned int    n_labels,
        unsigned int    n_images,
        unsigned int    n)
{
  for (unsigned int i = 0; i < n; ++i)
    { 
      x[i].worker = (lrand48 () >> 16) % n_workers;
      x[i].image = (lrand48 () >> 16) % n_images;
      x[i].rating = (lrand48 () >> 16) % n_labels;
    }
}

static void
test_to_from_identity_once (void)
{ 
  /* (kappa, logalpha, tau, gamma) = 
     from_x (kappa, logalpha, tau, gamma, to_x (kappa, logalpha, tau, gamma)) */

  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;

  double kappa;
  double logalpha[n_images];
  double tau[n_workers * (n_labels - 1)];
  double gamma[n_labels - 1];

  double newkappa;
  double newlogalpha[n_images];
  double newtau[n_workers * (n_labels - 1)];
  double newgamma[n_labels - 1];

  kappa = drand48 ();
  fill (logalpha, drand48, n_images);
  fill (tau, drand48, n_workers * (n_labels - 1));
  fill (gamma, drand48, n_labels - 1);

  newkappa = drand48 ();
  fill (newlogalpha, drand48, n_images);
  fill (newtau, drand48, n_workers * (n_labels - 1));
  fill (newgamma, drand48, n_labels - 1);

  size_t n_dims = 1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1);
  gsl_vector* x = gsl_vector_alloc (n_dims);

  to_x (&kappa, logalpha, n_workers, n_labels, tau, n_images, gamma, x);
  from_x (&newkappa, newlogalpha, n_workers, n_labels, newtau, n_images, newgamma, x);
          
  gsl_vector_free (x);

  assert (kappa == newkappa);
  assert (memcmp (logalpha, newlogalpha, n_images * sizeof (logalpha[0])) == 0);
  assert (memcmp (tau,
                  newtau,
                  n_workers * (n_labels - 1) * sizeof (tau[0])) == 0);
  assert (memcmp (gamma, newgamma, (n_labels - 1) * sizeof (gamma[0])) == 0);
}

static void
test_to_from_inverse (void)
{
  for (unsigned int i = 0; i < 100; ++i)
    { 
      test_to_from_identity_once ();
    }
}

static void
test_qfunc_f_once (void)
{
  /* qfunc (kappa, logalpha, tau, gamma) = 
     -qfunc_f (to_x (kappa, logalpha, tau, gamma)) */

  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_ratings = ((lrand48 () >> 16) % 5) + 1;

  double kappa;
  double logalpha[n_images];
  double tau[n_workers * (n_labels - 1)];
  double gamma[n_labels - 1];
  double pz[n_images * n_labels];
  double priorz[n_labels];
  Rating ratings[n_ratings];
  NormalDistribution priorkappa = normal_distribution (1.0 + drand48 (),
                                                       1.0 + drand48 ());
  ParameterizedNormalDistribution priorlogalpha = 
    parameterized_normal_distribution (1.0 + drand48 ());
  ParameterizedNormalDistribution priortau = 
    parameterized_normal_distribution (1.0 + drand48 ());
  NormalDistribution normalpriorgamma[n_labels - 1];
  Distribution* priorgamma[n_labels - 1];

  for (unsigned int i = 0; i + 1 < n_labels; ++i)
    {
      normalpriorgamma[i] = normal_distribution (1.0 + drand48 (),
                                                 1.0 + drand48 ());
      priorgamma[i] = &normalpriorgamma[i].base;
    }

  kappa = drand48 ();
  fill (logalpha, drand48, n_images);
  fill (tau, drand48, n_workers * (n_labels - 1));
  fill (gamma, drand48, n_labels - 1);
  fill (pz, drand48, n_images * n_labels);
  fill (priorz, drand48, n_labels);
  fill_r (ratings, n_workers, n_labels, n_images, n_ratings);

  QfuncParams qfunc_params = { .n_workers = n_workers,
                               .n_labels = n_labels,
                               .n_images = n_images,
                               .pz = pz,
                               .priorz = priorz,
                               .ratings = ratings,
                               .n_ratings = n_ratings,
                               .priorkappa = &priorkappa.base,
                               .priorlogalpha = &priorlogalpha.base,
                               .priortau = &priortau.base,
                               .priorgamma = priorgamma };

  size_t n_dims = 1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1);
  gsl_vector* x = gsl_vector_alloc (n_dims);

  to_x (&kappa, logalpha, n_workers, n_labels, tau, n_images, gamma, x);

  double q = qfunc (kappa, logalpha, n_workers, n_labels,
                    tau, n_images, gamma, pz, priorz,
                    ratings, n_ratings, &priorkappa.base,
                    &priorlogalpha.base, &priortau.base,
                    priorgamma);

  double qf = qfunc_f (x, &qfunc_params);

  assert (fabs (q + qf) <= 1e-8 * (1 + fabs (q) + fabs (qf)) ||
          (fprintf (stderr, "%g ?= %g\n", q, qf), 0));

  gsl_vector_free (x);
}

static void
test_qfunc_f (void)
{ 
  for (unsigned int i = 0; i < 100; ++i)
    { 
      test_qfunc_f_once ();
    }
}

static void
test_qfunc_df_once (void)
{
  /* dqfunc (kappa, logalpha, tau, gamma) = 
     -qfunc_df (to_x (kappa, logalpha, tau, gamma)) */

  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_ratings = ((lrand48 () >> 16) % 5) + 1;

  double kappa;
  double logalpha[n_images];
  double tau[n_workers * (n_labels - 1)];
  double gamma[n_labels - 1];
  double dkappa;
  double dlogalpha[n_images];
  double dtau[n_workers * (n_labels - 1)];
  double dgamma[n_labels - 1];
  double newdkappa;
  double newdlogalpha[n_images];
  double newdtau[n_workers * (n_labels - 1)];
  double newdgamma[n_labels - 1];
  double pz[n_images * n_labels];
  double priorz[n_labels];
  Rating ratings[n_ratings];
  NormalDistribution priorkappa = normal_distribution (1.0 + drand48 (),
                                                       1.0 + drand48 ());
  ParameterizedNormalDistribution priorlogalpha = 
    parameterized_normal_distribution (1.0 + drand48 ());
  ParameterizedNormalDistribution priortau = 
    parameterized_normal_distribution (1.0 + drand48 ());
  NormalDistribution normalpriorgamma[n_labels - 1];
  Distribution* priorgamma[n_labels - 1];

  for (unsigned int i = 0; i + 1 < n_labels; ++i)
    {
      normalpriorgamma[i] = normal_distribution (1.0 + drand48 (),
                                                 1.0 + drand48 ());
      priorgamma[i] = &normalpriorgamma[i].base;
    }

  kappa = drand48 ();
  fill (logalpha, drand48, n_images);
  fill (tau, drand48, n_workers * (n_labels - 1));
  fill (gamma, drand48, n_labels - 1);
  dkappa = drand48 ();
  fill (dlogalpha, drand48, n_images);
  fill (dtau, drand48, n_workers * (n_labels - 1));
  fill (dgamma, drand48, n_labels - 1);
  newdkappa = drand48 ();
  fill (newdlogalpha, drand48, n_images);
  fill (newdtau, drand48, n_workers * (n_labels - 1));
  fill (newdgamma, drand48, n_labels - 1);
  fill (pz, drand48, n_images * n_labels);
  fill (priorz, drand48, n_labels);
  fill_r (ratings, n_workers, n_labels, n_images, n_ratings);

  QfuncParams qfunc_params = { .n_workers = n_workers,
                               .n_labels = n_labels,
                               .n_images = n_images,
                               .pz = pz,
                               .priorz = priorz,
                               .ratings = ratings,
                               .n_ratings = n_ratings,
                               .priorkappa = &priorkappa.base,
                               .priorlogalpha = &priorlogalpha.base,
                               .priortau = &priortau.base,
                               .priorgamma = priorgamma };

  size_t n_dims = 1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1);
  gsl_vector* x = gsl_vector_alloc (n_dims);
  gsl_vector* g = gsl_vector_alloc (n_dims);

  to_x (&kappa, logalpha, n_workers, n_labels, tau, n_images, gamma, x);

  dqfunc (&dkappa, dlogalpha, dtau, dgamma,
          kappa, logalpha, n_workers, n_labels,
          tau, n_images, gamma, pz,
          ratings, n_ratings, &priorkappa.base,
          &priorlogalpha.base, &priortau.base,
          priorgamma);

  qfunc_df (x, &qfunc_params, g);

  from_x (&newdkappa, newdlogalpha, n_workers, n_labels, newdtau, n_images, newdgamma, g);

  assert (fabs (dkappa + newdkappa) <= 
          1e-8 * (1 + fabs (dkappa) + fabs (newdkappa)) ||
          (fprintf (stderr, "%g ?= %g\n", dkappa, newdkappa), 0));

  for (unsigned int i = 0; i < n_images; ++i)
    {
      assert (fabs (dlogalpha[i] + newdlogalpha[i]) <=
              1e-8 * (1 + fabs (dlogalpha[i]) + fabs (newdlogalpha[i])) ||
              (fprintf (stderr, "%g ?= %g\n", dlogalpha[i], newdlogalpha[i]), 0));
    }

  for (unsigned int i = 0; i < n_workers * (n_labels - 1); ++i)
    {
      assert (fabs (dtau[i] + newdtau[i]) <=
              1e-8 * (1 + fabs (dtau[i]) + fabs (newdtau[i])) ||
              (fprintf (stderr, "%g ?= %g\n", dtau[i], newdtau[i]), 0));
    }

  for (unsigned int i = 0; i < (n_labels - 1); ++i)
    {
      assert (fabs (dgamma[i] + newdgamma[i]) <=
              1e-8 * (1 + fabs (dgamma[i]) + fabs (newdgamma[i])) ||
              (fprintf (stderr, "%g ?= %g\n", dgamma[i], newdgamma[i]), 0));
    }

  gsl_vector_free (g);
  gsl_vector_free (x);
}

static void
test_qfunc_df (void)
{ 
  for (unsigned int i = 0; i < 100; ++i)
    { 
      test_qfunc_df_once ();
    }
}

static void
test_mstep_no_data_once (void)
{
  /* without data mstep should converge to prior maximum */

  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;

  double kappa;
  double logalpha[n_images];
  double tau[n_workers * (n_labels - 1)];
  double gamma[n_labels - 1];
  double pz[n_images * n_labels];
  double priorz[n_labels];
  NormalDistribution priorkappa = normal_distribution (1.0 + drand48 (),
                                                       1.0 + drand48 ());
  ParameterizedNormalDistribution priorlogalpha = 
    parameterized_normal_distribution (1.0 + drand48 ());
  ParameterizedNormalDistribution priortau = 
    parameterized_normal_distribution (1.0 + drand48 ());
  NormalDistribution normalpriorgamma[n_labels - 1];
  Distribution* priorgamma[n_labels - 1];

  for (unsigned int i = 0; i + 1 < n_labels; ++i)
    {
      normalpriorgamma[i] = normal_distribution (1.0 + drand48 (),
                                                 1.0 + drand48 ());
      priorgamma[i] = &normalpriorgamma[i].base;
    }

  kappa = drand48 ();
  fill (logalpha, drand48, n_images);
  fill (tau, drand48, n_workers * (n_labels - 1));
  fill (gamma, drand48, n_labels - 1);
  fill (pz, drand48, n_images * n_labels);
  fill (priorz, drand48, n_labels);

  mstep (&kappa, logalpha, n_workers, n_labels,
         tau, n_images, gamma, pz, priorz,
         NULL, 0, &priorkappa.base, &priorlogalpha.base,
         &priortau.base, priorgamma, 100, 1e-8);

  assert (   fabs (kappa - priorkappa.mean) 
          <= 1e-4 * (1 + fabs (kappa) + fabs (priorkappa.mean))
          || (fprintf (stderr, "%g ?= %g\n", kappa, priorkappa.mean), 0));

  for (unsigned int i = 0; i < n_images; ++i)
    {
      assert (   fabs (logalpha[i] - priorkappa.mean)
              <= 1e-4 * (1 + fabs (logalpha[i]) + fabs (priorkappa.mean))
              || (fprintf (stderr, "%u: %g ?= %g\n", i, logalpha[i], priorkappa.mean), 0));
    }

  for (unsigned int i = 0; i + 1 < n_labels; ++i)
    {
      assert (   fabs (gamma[i] - normalpriorgamma[i].mean)
              <= 1e-4 * (1 + fabs (gamma[i]) + fabs (normalpriorgamma[i].mean))
              || (fprintf (stderr, "%u: %g ?= %g\n", i, gamma[i], normalpriorgamma[i].mean), 0));
    }

  for (unsigned int i = 0; i < n_workers; ++i)
    {
      for (unsigned int j = 0; j + 1 < n_labels; ++j)
        {
          assert (   fabs (tau[i * (n_labels - 1) + j] - normalpriorgamma[j].mean)
                  <= 1e-4 * (1 + fabs (tau[i * (n_labels - 1) + j]) + fabs (normalpriorgamma[j].mean))
                  || (fprintf (stderr, "%u: %g ?= %g\n", i, tau[i * (n_labels - 1) + j], normalpriorgamma[j].mean), 0));
        }
    }
}

static void
test_mstep_no_data (void)
{
  for (unsigned int i = 0; i < 20; ++i)
    {
      test_mstep_no_data_once ();
    }
}

typedef struct _MstepTestDatum MstepTestDatum;
struct _MstepTestDatum 
{
  struct
    {
      double kappa;
      double logalpha[3];
      double tau[4*(5-1)];
      double gamma[(5-1)];
      double pz[3*5];
      double priorz[5];
      Rating ratings[8];
      double priorlogalphavar;
      double priortauvar;
      double priorgamma1mean;
      double priorgamma1var;
      double priorgamma2mean;
      double priorgamma2var;
      double priorgamma3mean;
      double priorgamma3var;
      double priorgamma4mean;
      double priorgamma4var;
      double priorkappamean;
      double priorkappavar;
    }                           input;

  struct
    {
      double q;
      struct
        {
          double kappa;
          double logalpha[3];
          double tau[4*(5-1)];
          double gamma[(5-1)];
        }    v;
    }                           output;
};

static void
test_mstep (void)
{
  MstepTestDatum data[] = 
    {{{0.513223, {0.370493, 0.0648764, 0.271536}, {0.0659061, 0.329566,  0.651285, 0.830363, 0.172057, 0.0864764, 0.492879, 0.502661,  0.665746, 0.916565, 0.415802, 0.655563, 0.879574, 0.555494,  0.697461, 0.220453}, {0.891685, 0.0711726, 0.859653,  0.289465}, {0.378462, 0.700679, 0.794777, 0.0179292, 0.312556,  0.371114, 0.143491, 0.187566, 0.140499, 0.284637, 0.650612,  0.684906, 0.474754, 0.368072, 0.23481}, {0.0293429, 0.595179,  0.812578, 0.537349,  0.80889}, {{1, 1, 4}, {2, 2, 2}, {0, 2, 4}, {0, 1, 0}, {3, 2,  3}, {3, 1, 3}, {2, 2, 3}, {0, 1, 2}}, 0.703494, 0.741405, 1.6777,  1.51942, 1.32503, 1.04073, 1.88292, 1.5015, 1.01248, 1.66961,  1.73943,  1.31393}, {-49.787796763029122703, {-0.80921647278382306224, {-0.80921647280419557748, -1.4218376123998896279, -0.92720785323178537257}, {0.86028690772310295095,  0.18232607821975874039, 0.38243449696575770786,  1.4040220369689717588, 0.80213026154320524117,  0.20890868834181002626, 0.37341780648193863964,  1.5638551213012271572,  0.68634455253279558992, -0.019945648184538371304,  0.49829518136473487928, 1.9636466175464791776,  0.73626693252703082012, 0.079782511682622066166,  0.15426623794029290286,  1.8687992258024993041}, {0.82218095310869945015,  0.24925758193466393513, 0.44005218882713505078,  1.6677764244257724710}}}}, {{0.871976, {0.384975, 0.0888163,  0.629023}, {0.397223, 0.0169028, 0.854006, 0.59968, 0.802043,  0.204325, 0.316657, 0.79079, 0.0985488, 0.46292, 0.638961,  0.271366, 0.773517, 0.422194, 0.756041, 0.769871}, {0.761041,  0.752583, 0.0166129, 0.455942}, {0.889065, 0.367608, 0.927797,  0.82692, 0.491842, 0.350705, 0.0737907, 0.22724, 0.689799,  0.14638, 0.757134, 0.43645, 0.59125, 0.68346,  0.118173}, {0.165084, 0.817733, 0.261265, 0.362132,  0.395212}, {{0, 1, 1}, {0, 1, 0}, {3, 2, 2}, {3, 1, 1}, {0, 1,  1}, {3, 0, 4}, {0, 0, 4}, {0, 0, 1}}, 0.056692, 0.508683,  1.34552, 1.93927, 1.16763, 1.14107, 1.41772, 1.11235, 1.67579,  1.79037, 1.34393,  1.88511}, {-46.507243283941960336, {-1.9563687911572665421, {-1.9548259789166271928, -1.9603872904975462149, -1.9568779640363750149}, {0.45548481225305266324,  1.8017108848005011208, 1.7594701824280544091,  0.65042152378511438280, 0.46971983475123617154,  1.6591621048076775692, 1.6971790911146968311,  0.67218922082964113161, 0.46971983464509730788,  1.6591621048466927674, 1.6971790911633240390,  0.67218922075127363996, 0.42369582706736403481,  1.6142965762962619161, 1.6933299853379904969,  0.61294154614155021546}, {0.46971983468242881861,  1.6591621048551097666, 1.6971790911869478309,  0.67218922085077237428}}}}, {{0.985986, {0.64399, 0.586798,  0.448661}, {0.394736, 0.96053, 0.468625, 0.283577, 0.577003,  0.699265, 0.106493, 0.888365, 0.520311, 0.190582, 0.760973,  0.949095, 0.352684, 0.0495071, 0.343251, 0.836745}, {0.676899,  0.259137, 0.999319, 0.951634}, {0.690912, 0.615147, 0.412521,  0.502973, 0.296176, 0.654617, 0.943896, 0.219396, 0.719173,  0.955353, 0.837403, 0.331031, 0.198863, 0.764771,  0.07643}, {0.381936, 0.846179, 0.715264, 0.733179,  0.545191}, {{3, 2, 0}, {1, 2, 4}, {1, 0, 2}, {2, 1, 0}, {2, 2,  1}, {2, 2, 4}, {2, 1, 1}, {2, 0, 1}}, 0.16928, 0.456126, 1.73386,  1.59356, 1.47837, 1.84098, 1.32134, 1.09058, 1.18219, 1.18636,  1.37744,  1.87119}, {-46.140267779628954328, {-1.9603966639549733313, {-1.9705982479921820832, -1.9631949121519885066, -1.9747145264294106337}, {1.9701481380088261656,  3.4925403651021280247, 1.8596230608079246105,  1.1254217528792908809, 1.9364489445341498150,  3.4282525394488741624, 1.8450005889184871450,  1.0877869494329756288, 1.9756805757527048446,  3.6474031920712986131, 1.9461377996416939508,  1.1429088218881008322, 2.0176736066407139148,  3.5256083716743816919, 1.8818901510239950300,  1.1371776796913919943}, {1.9701481379851129585,  3.4925403651424497644, 1.8596230609055339343,  1.1254217528098298846}}}}};

  for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
    {
      ParameterizedNormalDistribution priorlogalpha = 
        parameterized_normal_distribution (data[i].input.priorlogalphavar);

      ParameterizedNormalDistribution priortau =
        parameterized_normal_distribution (data[i].input.priortauvar);

      NormalDistribution priorgammanormal[4] = {
          normal_distribution (data[i].input.priorgamma1mean,
                               data[i].input.priorgamma1var),
          normal_distribution (data[i].input.priorgamma2mean,
                               data[i].input.priorgamma2var),
          normal_distribution (data[i].input.priorgamma3mean,
                               data[i].input.priorgamma3var),
          normal_distribution (data[i].input.priorgamma4mean,
                               data[i].input.priorgamma4var)
      };

      Distribution* priorgamma[4] = {
          &priorgammanormal[0].base,
          &priorgammanormal[1].base,
          &priorgammanormal[2].base,
          &priorgammanormal[3].base
      };

      NormalDistribution priorkappa = 
        normal_distribution (data[i].input.priorkappamean,
                             data[i].input.priorkappavar);


      double kappa = data[i].input.kappa;
      double logalpha[3];
      double tau[4*(5-1)];
      double gamma[(5-1)];

      memcpy (logalpha, data[i].input.logalpha, 3 * sizeof (double));
      memcpy (tau, data[i].input.tau, 4*(5-1) * sizeof (double));
      memcpy (gamma, data[i].input.gamma, (5-1) * sizeof (double));

      double q = mstep (&kappa,
                        logalpha,
                        4,
                        5,
                        tau,
                        3,
                        gamma,
                        data[i].input.pz,
                        data[i].input.priorz,
                        data[i].input.ratings,
                        8,
                        &priorkappa.base,
                        &priorlogalpha.base,
                        &priortau.base,
                        priorgamma,
                        10000,
                        1e-16);

      assert (      fabs (q - data[i].output.q) 
                 <= 1e-4 * (1 + fabs (q) + fabs (data[i].output.q))
              || (fprintf (stderr, "fail: %g != %g\n", q, data[i].output.q),
                  0));

      assert (   fabs (kappa - data[i].output.v.kappa)
              <= 1e-4 * (1 + fabs (kappa) + fabs (data[i].output.v.kappa))
              || (fprintf (stderr, "%g ?= %g\n", kappa, data[i].output.v.kappa), 0));

      for (unsigned int j = 0; j < 3; ++j)
        {
          assert (   fabs (logalpha[j] - data[i].output.v.logalpha[j])
                  <= 1e-4 * (1 + fabs (logalpha[j]) + fabs (data[i].output.v.logalpha[j]))
                  || (fprintf (stderr, "%u: %g ?= %g\n", j, logalpha[j], data[i].output.v.logalpha[j]), 0));
        }

      for (unsigned int j = 0; j < 4; ++j)
        {
          for (unsigned int k = 0; k + 1 < 5; ++k)
            {
              assert (   fabs (tau[j * (5 - 1) + k] - data[i].output.v.tau[j * (5 - 1) + k])
                      <= 1e-4 * (1 + fabs (tau[j * (5 - 1) + k]) + fabs (data[i].output.v.tau[j * (5 - 1) + k]))
                      || (fprintf (stderr, "%u: %g ?= %g\n", j, tau[j * (5 - 1) + k], data[i].output.v.tau[j * (5 - 1) + k]), 0));
            }
        }

      for (unsigned int j = 0; j + 1 < 5; ++j)
        {
          assert (   fabs (gamma[j] - data[i].output.v.gamma[j])
                  <= 1e-4 * (1 + fabs (gamma[j]) + fabs (data[i].output.v.gamma[j]))
                  || (fprintf (stderr, "%u: %g ?= %g\n", j, gamma[j], data[i].output.v.gamma[j]), 0));
        }

    }
}

int
main (void)
{
  srand48 (69);

  test_to_from_inverse ();
  test_qfunc_f ();
  test_qfunc_df ();
  test_mstep_no_data ();
  test_mstep ();

  return 0;
}
